from typing import List
import pathlib
import re
import zipfile
from copy import deepcopy
import sys

from lxml import etree

from src.students import read_students,Student

namespaces = {
    'wpc':"http://schemas.microsoft.com/office/word/2010/wordprocessingCanvas",
    'mc':"http://schemas.openxmlformats.org/markup-compatibility/2006",
    'o':"urn:schemas-microsoft-com:office:office",
    'r':"http://schemas.openxmlformats.org/officeDocument/2006/relationships",
    'm':"http://schemas.openxmlformats.org/officeDocument/2006/math",
    'v':"urn:schemas-microsoft-com:vml",
    'wp14':"http://schemas.microsoft.com/office/word/2010/wordprocessingDrawing",
    'wp':"http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
    'w':"http://schemas.openxmlformats.org/wordprocessingml/2006/main",
    'w14':"http://schemas.microsoft.com/office/word/2010/wordml",
    'w10':"urn:schemas-microsoft-com:office:word",
    'w15':"http://schemas.microsoft.com/office/word/2012/wordml",
    'wpg':"http://schemas.microsoft.com/office/word/2010/wordprocessingGroup",
    'wpi':"http://schemas.microsoft.com/office/word/2010/wordprocessingInk",
    'wne':"http://schemas.microsoft.com/office/word/2006/wordml",
    'wps':"http://schemas.microsoft.com/office/word/2010/wordprocessingShape",
    'wpsCustomData':"http://www.wps.cn/officeDocument/2013/wpsCustomData",
}

def get_problem_text_to_nth_element(elem: etree._Element, n: int) -> str | None: 
    text_nodes = elem.xpath(".//w:t", namespaces = namespaces)
    if len(text_nodes) == 0:
        return None

    actual_n = min(n, len(text_nodes))

    text_node = text_nodes[actual_n - 1]
    return "".join(text_node.itertext())

problem_number_regex = re.compile(r'^(\d+)．');
def get_problem_number(elem: etree._Element) -> int | None:
    text = get_problem_text_to_nth_element(elem, 2)
    if text is None:
        return None

    matched = problem_number_regex.match(text)
    if matched is None:
        return None

    return int(matched.group(1))

def build_problem_set(doc: etree._Element) -> dict[int, List[etree._Element]]:
    result = dict();
    current_problem_number:int | None = None

    for orig_para in doc.xpath('/w:document/w:body/w:p', namespaces=namespaces):
        problem_number = get_problem_number(orig_para)

        if problem_number is None:
            if current_problem_number is not None:
                result[current_problem_number].append(orig_para)
            continue

        if result.get(problem_number) is not None:
            raise Exception(f'Duplicate problem number problem_number')

        current_problem_number = problem_number
        result[problem_number] = [deepcopy(orig_para)]

    return result

def student_title(student: Student, problem_set_name: str):
    elem = etree.Element(f'{{{namespaces["w"]}}}p', nsmap=namespaces)
    run = etree.SubElement(elem, f'{{{namespaces["w"]}}}r')
    run_props = etree.SubElement(run, f'{{{namespaces["w"]}}}rPr')
    font = etree.SubElement(run_props, f'{{{namespaces["w"]}}}rFonts')
    font.set(f'{{{namespaces["w"]}}}hint', 'eastAsia')
    font.set(f'{{{namespaces["w"]}}}ascii', '宋体')
    font.set(f'{{{namespaces["w"]}}}hAnsi', '宋体')
    font.set(f'{{{namespaces["w"]}}}cs', '宋体')
    size = etree.SubElement(run_props, f'{{{namespaces["w"]}}}sz')
    size.set(f'{{{namespaces["w"]}}}val', '21')
    etree.SubElement(run_props, f'{{{namespaces["w"]}}}b')
    text = etree.SubElement(run, f'{{{namespaces["w"]}}}t')
    text.text = f'{student.number} {student.name} {problem_set_name}' # type: ignore
    return elem

def student_margin_end(): 
    elem = etree.Element(f'{{{namespaces["w"]}}}p', nsmap=namespaces)
    para_props = etree.SubElement(elem, f'{{{namespaces["w"]}}}pPr')
    spacing = etree.SubElement(para_props, f'{{{namespaces["w"]}}}spacing')
    spacing.set(f'{{{namespaces["w"]}}}line', '360')
    return elem

def section_para():
    elem = etree.Element(f'{{{namespaces["w"]}}}p', nsmap=namespaces)
    para_props = etree.SubElement(elem, f'{{{namespaces["w"]}}}pPr')
    sect_pr = etree.SubElement(para_props, f'{{{namespaces["w"]}}}sectPr')
    page_size = etree.SubElement(sect_pr, f'{{{namespaces["w"]}}}pgSz')
    page_size.set(f'{{{namespaces["w"]}}}w', '11907')
    page_size.set(f'{{{namespaces["w"]}}}h', '16839')
    page_margin = etree.SubElement(sect_pr, f'{{{namespaces["w"]}}}pgMar')
    page_margin.set(f'{{{namespaces["w"]}}}top', '180')
    page_margin.set(f'{{{namespaces["w"]}}}right', '227')
    page_margin.set(f'{{{namespaces["w"]}}}bottom', '479')
    page_margin.set(f'{{{namespaces["w"]}}}left', '340')
    page_margin.set(f'{{{namespaces["w"]}}}header', '851')
    page_margin.set(f'{{{namespaces["w"]}}}footer', '425')
    page_margin.set(f'{{{namespaces["w"]}}}gutter', '0')
    cols = etree.SubElement(sect_pr, f'{{{namespaces["w"]}}}cols')
    cols.set(f'{{{namespaces["w"]}}}space', '427')
    cols.set(f'{{{namespaces["w"]}}}num', '2')
    doc_grid = etree.SubElement(sect_pr, f'{{{namespaces["w"]}}}docGrid')
    doc_grid.set(f'{{{namespaces["w"]}}}type', 'lines')
    doc_grid.set(f'{{{namespaces["w"]}}}linePitch', '312')
    doc_grid.set(f'{{{namespaces["w"]}}}charSpace', '0')
    return elem

if len(sys.argv) < 2:
    raise Exception("missing title argument")
title = sys.argv[1]
assert title is not None

base_dir = pathlib.Path(__file__).resolve().parent.parent
local_dir = base_dir / 'local'
b_doc_path = local_dir / 'B.docx'
a_xlsx_path = local_dir / 'A.xlsx'
out_doc_path = local_dir / 'out.docx'

students = read_students(str(a_xlsx_path))

with zipfile.ZipFile(b_doc_path, 'r') as b_doc_zip:
    with zipfile.ZipFile(out_doc_path, 'w') as out_doc_zip:
        for file in b_doc_zip.filelist:
            if file.filename == 'word/document.xml':
                b_main = etree.fromstring(b_doc_zip.read('word/document.xml'))
                problem_set = build_problem_set(b_main)
                new_doc = deepcopy(b_main)
                new_body = new_doc.find('.//w:body', namespaces = namespaces)
                if new_body is None:
                    raise Exception()
                new_body.clear()

                # for problem_para in problem_set.values():
                #     for para in problem_para:
                #         new_body.append(deepcopy(para))

                for student in students:
                    new_body.append(student_title(student, title))
                    for problem_id in sorted(student.failed_problems):
                        problem_paras = problem_set.get(problem_id)
                        assert problem_paras is not None
                        for para in problem_paras:
                            new_body.append(deepcopy(para))
                    new_body.append(student_margin_end())

                new_body.append(section_para())

                out_doc_zip.writestr('word/document.xml', etree.tostring(new_doc, encoding="utf-8", pretty_print=False, with_tail=False))
            else:
                out_doc_zip.writestr(file, b_doc_zip.read(file))

